import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.linalg import solve_continuous_are

class WheeledLegRobot:
    def __init__(self):
        # Robot parameters
        self.m_body = 2.0      # body mass kg
        self.m_leg = 0.5       # leg mass kg
        self.l_hip = 0.3       # hip joint length m
        self.l_knee = 0.3      # knee joint length m
        self.I_body = 0.1      # body moment of inertia
        self.g = 9.81          # gravity
        
        # State: [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot]
        self.state_dim = 12
        self.control_dim = 6
        
    def dynamics(self, state, control_torques, t, disturbance_type='none'):
        """Robot dynamics equations with disturbances"""
        x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot = state
        tau1, tau2, tau3, tau4, tau5, tau6 = control_torques
        
        # Simplified dynamics model
        M = np.diag([self.m_body, self.m_body, self.m_body, 
                     self.I_body, self.I_body, self.I_body])
        
        # Coriolis and gravity terms
        C = np.array([0, 0, self.m_body * self.g, 0, 0, 0])
        
        # Control input mapping
        B = np.eye(6)
        
        # Add disturbances based on type
        disturbance = self.get_disturbance(t, disturbance_type)
        
        # Acceleration calculation with disturbance
        accelerations = np.linalg.solve(M, B @ control_torques - C + disturbance)
        
        # State derivatives
        state_dot = np.zeros(12)
        state_dot[0:6] = state[6:12]  # position derivatives = velocities
        state_dot[6:12] = accelerations  # velocity derivatives = accelerations
        
        return state_dot
    
    def get_disturbance(self, t, disturbance_type):
        """Generate different types of disturbances"""
        disturbance = np.zeros(6)
        
        if disturbance_type == 'none':
            return disturbance
            
        elif disturbance_type == 'constant':
            # Constant disturbance force
            disturbance = np.array([2.0, 1.5, 3.0, 0.5, 0.5, 0.5])
            
        elif disturbance_type == 'sinusoidal':
            # Sinusoidal disturbance
            amplitude = 3.0
            frequency = 2.0
            disturbance = amplitude * np.array([
                np.sin(2*np.pi*frequency*t),
                np.cos(2*np.pi*frequency*t),
                0.5*np.sin(2*np.pi*frequency*t + np.pi/4),
                0.3*np.cos(2*np.pi*frequency*t),
                0.3*np.sin(2*np.pi*frequency*t),
                0.2*np.cos(2*np.pi*frequency*t)
            ])
            
        elif disturbance_type == 'impulse':
            # Impulse disturbances at specific times
            if 2.0 <= t < 2.1:
                disturbance = np.array([10, 8, 5, 2, 2, 2])
            elif 5.0 <= t < 5.1:
                disturbance = np.array([-8, 10, 3, 1, 1, 1])
            else:
                disturbance = np.zeros(6)
                
        elif disturbance_type == 'random':
            # Random noise disturbance
            disturbance = 2.0 * np.random.randn(6)
            
        elif disturbance_type == 'step':
            # Step disturbance
            if t >= 3.0:
                disturbance = np.array([4.0, 3.0, 2.0, 1.0, 1.0, 1.0])
                
        elif disturbance_type == 'complex':
            # Complex combined disturbance
            if t < 3.0:
                # Constant
                disturbance = np.array([1.0, 0.5, 2.0, 0.3, 0.3, 0.3])
            elif 3.0 <= t < 6.0:
                # Sinusoidal
                amplitude = 2.0
                disturbance = amplitude * np.array([
                    np.sin(2*np.pi*1.5*t),
                    np.cos(2*np.pi*1.5*t),
                    0,
                    0.2*np.sin(2*np.pi*1.5*t),
                    0.2*np.cos(2*np.pi*1.5*t),
                    0
                ])
            else:
                # Random
                disturbance = 1.5 * np.random.randn(6)
                
        return disturbance

class PIDController:
    def __init__(self, kp, ki, kd, dt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.dt = dt
        self.integral = np.zeros(6)
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        self.integral += error * self.dt
        
        # Anti-windup
        self.integral = np.clip(self.integral, -50, 50)
        
        control = (self.kp * error + 
                  self.ki * self.integral + 
                  self.kd * error_derivative)
        
        self.previous_error = error.copy()
        return control

class SMCController:
    def __init__(self, lambda_val, k, phi, dt):
        self.lambda_val = lambda_val
        self.k = k
        self.phi = phi
        self.dt = dt
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        # Sliding surface
        s = error_derivative + self.lambda_val * error
        
        # Control law (using saturation function to reduce chattering)
        sat_s = np.tanh(s * 5)  # Reduced gain for better robustness
        control = self.k * sat_s + self.phi * s
        
        self.previous_error = error.copy()
        return control

class LQRController:
    def __init__(self, A, B, Q, R):
        # Use scipy to solve continuous-time algebraic Riccati equation
        P = solve_continuous_are(A, B, Q, R)
        self.K = np.linalg.solve(R, B.T @ P)
    
    def compute_control(self, desired_state, current_state):
        error = current_state - desired_state
        return -self.K @ error

def simulate_controller_with_disturbance(controller, robot, desired_trajectory, t_span, dt, disturbance_type):
    """Simulate specific controller with disturbance"""
    t_eval = np.arange(t_span[0], t_span[1], dt)
    n_steps = len(t_eval)
    
    # Initialize state records
    states = np.zeros((n_steps, robot.state_dim))
    controls = np.zeros((n_steps, robot.control_dim))
    errors = np.zeros((n_steps, 6))
    disturbances = np.zeros((n_steps, 6))
    
    # Initial state
    state = np.array([0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    
    for i, t in enumerate(t_eval):
        # Get desired state
        desired_state = desired_trajectory(t)
        
        # Calculate control input
        control = controller.compute_control(desired_state, state)
        
        # Limit control input
        control = np.clip(control, -25, 25)  # Increased limits for disturbance rejection
        
        # Get current disturbance
        current_disturbance = robot.get_disturbance(t, disturbance_type)
        
        # Record state and control
        states[i] = state
        controls[i] = control
        errors[i] = desired_state[0:6] - state[0:6]
        disturbances[i] = current_disturbance
        
        # Integrate dynamics with disturbance
        state_dot = robot.dynamics(state, control, t, disturbance_type)
        state = state + state_dot * dt
    
    return t_eval, states, controls, errors, disturbances

def desired_circular_trajectory(t):
    """Circular trajectory"""
    radius = 1.0
    omega = 0.5
    
    x = radius * np.cos(omega * t)
    y = radius * np.sin(omega * t)
    z = 0.5  # maintain fixed height
    theta, phi, psi = 0, 0, 0
    x_dot = -radius * omega * np.sin(omega * t)
    y_dot = radius * omega * np.cos(omega * t)
    z_dot = 0
    theta_dot, phi_dot, psi_dot = 0, 0, 0
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def run_robustness_test(disturbance_type):
    """Run robustness test for all controllers with specified disturbance"""
    print(f"\n{'='*60}")
    print(f"ROBUSTNESS TEST: {disturbance_type.upper()} DISTURBANCE")
    print(f"{'='*60}")
    
    # Initialize robot and controllers
    robot = WheeledLegRobot()
    dt = 0.01
    t_span = (0, 10)
    
    # Controllers with slightly tuned parameters for robustness
    pid = PIDController(
        kp=np.array([18, 18, 22, 12, 12, 12]),  # Increased gains
        ki=np.array([0.2, 0.2, 0.8, 0.2, 0.2, 0.2]),
        kd=np.array([6, 6, 10, 4, 4, 4]),
        dt=dt
    )
    
    smc = SMCController(
        lambda_val=3.0,  # Increased for faster convergence
        k=np.array([12, 12, 15, 8, 8, 8]),  # Increased for disturbance rejection
        phi=np.array([3, 3, 4, 2, 2, 2]),
        dt=dt
    )
    
    # LQR Controller
    A = np.zeros((12, 12))
    A[0:6, 6:12] = np.eye(6)
    A[8, 8] = -0.2  # Increased damping
    
    B = np.zeros((12, 6))
    B[6:12, 0:6] = np.eye(6) / robot.m_body
    B[9:12, 3:6] = np.eye(3) / robot.I_body
    
    Q = np.diag([15, 15, 25, 8, 8, 8, 2, 2, 3, 1, 1, 1])  # Increased state penalties
    R = np.diag([0.05, 0.05, 0.05, 0.05, 0.05, 0.05])  # Reduced control penalty
    
    lqr = LQRController(A, B, Q, R)
    
    # Run simulations with disturbance
    print("Running simulations with disturbance...")
    t_pid, states_pid, controls_pid, errors_pid, dist_pid = simulate_controller_with_disturbance(
        pid, robot, desired_circular_trajectory, t_span, dt, disturbance_type)
    
    t_smc, states_smc, controls_smc, errors_smc, dist_smc = simulate_controller_with_disturbance(
        smc, robot, desired_circular_trajectory, t_span, dt, disturbance_type)
    
    t_lqr, states_lqr, controls_lqr, errors_lqr, dist_lqr = simulate_controller_with_disturbance(
        lqr, robot, desired_circular_trajectory, t_span, dt, disturbance_type)
    
    # Plot results
    plot_robustness_comparison(t_pid, t_smc, t_lqr, 
                             states_pid, states_smc, states_lqr,
                             errors_pid, errors_smc, errors_lqr,
                             controls_pid, controls_smc, controls_lqr,
                             dist_pid, disturbance_type)
    
    return calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, 
                                      controls_pid, controls_smc, controls_lqr,
                                      disturbance_type)

def calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, 
                               controls_pid, controls_smc, controls_lqr,
                               disturbance_type):
    """Calculate robustness performance metrics"""
    
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    control_norm_pid = np.linalg.norm(controls_pid, axis=1)
    control_norm_smc = np.linalg.norm(controls_smc, axis=1)
    control_norm_lqr = np.linalg.norm(controls_lqr, axis=1)
    
    # Recovery time after major disturbances
    def calculate_recovery_time(error_norm, threshold=0.1):
        max_error_idx = np.argmax(error_norm)
        for i in range(max_error_idx, len(error_norm)):
            if error_norm[i] < threshold and np.all(error_norm[i:i+100] < threshold):
                return i * 0.01  # Convert to seconds
        return len(error_norm) * 0.01
    
    # Steady-state error (last 2 seconds average)
    steady_start_idx = int(0.8 * len(error_norm_pid))  # Last 20% of simulation
    steady_state_pid = np.mean(error_norm_pid[steady_start_idx:])
    steady_state_smc = np.mean(error_norm_smc[steady_start_idx:])
    steady_state_lqr = np.mean(error_norm_lqr[steady_start_idx:])
    
    # Control effort metrics
    mean_control_pid = np.mean(control_norm_pid)
    mean_control_smc = np.mean(control_norm_smc)
    mean_control_lqr = np.mean(control_norm_lqr)
    
    max_control_pid = np.max(control_norm_pid)
    max_control_smc = np.max(control_norm_smc)
    max_control_lqr = np.max(control_norm_lqr)
    
    metrics = {
        'RMSE Position': [
            np.sqrt(np.mean(error_norm_pid**2)),
            np.sqrt(np.mean(error_norm_smc**2)),
            np.sqrt(np.mean(error_norm_lqr**2))
        ],
        'Max Position Error': [
            np.max(error_norm_pid),
            np.max(error_norm_smc),
            np.max(error_norm_lqr)
        ],
        'Steady-State Error': [
            steady_state_pid,
            steady_state_smc,
            steady_state_lqr
        ],
        'Recovery Time (s)': [
            calculate_recovery_time(error_norm_pid),
            calculate_recovery_time(error_norm_smc),
            calculate_recovery_time(error_norm_lqr)
        ],
        'Mean Control Effort': [
            mean_control_pid,
            mean_control_smc,
            mean_control_lqr
        ],
        'Max Control Effort': [
            max_control_pid,
            max_control_smc,
            max_control_lqr
        ]
    }
    
    return metrics

def plot_robustness_comparison(t_pid, t_smc, t_lqr, states_pid, states_smc, states_lqr,
                             errors_pid, errors_smc, errors_lqr, controls_pid, controls_smc, controls_lqr,
                             disturbances, disturbance_type):
    """Plot robustness comparison results"""
    
    fig, axes = plt.subplots(3, 3, figsize=(16, 12))
    fig.suptitle(f'Robustness Test: {disturbance_type.upper()} Disturbance', fontsize=16, fontweight='bold')
    
    # Position tracking with disturbance indication
    axes[0,0].plot(t_pid, states_pid[:, 0], 'r-', label='PID', linewidth=2)
    axes[0,0].plot(t_smc, states_smc[:, 0], 'g-', label='SMC', linewidth=2)
    axes[0,0].plot(t_lqr, states_lqr[:, 0], 'b-', label='LQR', linewidth=2)
    
    # Plot desired trajectory
    t_ref = np.linspace(0, 10, 100)
    x_ref = np.cos(0.5 * t_ref)
    axes[0,0].plot(t_ref, x_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    
    # Shade disturbance periods
    if disturbance_type in ['impulse', 'step']:
        axes[0,0].axvspan(2.0, 2.1, alpha=0.2, color='red', label='Disturbance')
        if disturbance_type == 'impulse':
            axes[0,0].axvspan(5.0, 5.1, alpha=0.2, color='red')
        elif disturbance_type == 'step':
            axes[0,0].axvspan(3.0, 10, alpha=0.2, color='red')
    
    axes[0,0].set_ylabel('X Position (m)')
    axes[0,0].set_title('X-Axis Position Tracking')
    axes[0,0].legend()
    axes[0,0].grid(True)
    
    # Y position
    axes[0,1].plot(t_pid, states_pid[:, 1], 'r-', label='PID', linewidth=2)
    axes[0,1].plot(t_smc, states_smc[:, 1], 'g-', label='SMC', linewidth=2)
    axes[0,1].plot(t_lqr, states_lqr[:, 1], 'b-', label='LQR', linewidth=2)
    
    y_ref = np.sin(0.5 * t_ref)
    axes[0,1].plot(t_ref, y_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,1].set_ylabel('Y Position (m)')
    axes[0,1].set_title('Y-Axis Position Tracking')
    axes[0,1].legend()
    axes[0,1].grid(True)
    
    # Z position
    axes[0,2].plot(t_pid, states_pid[:, 2], 'r-', label='PID', linewidth=2)
    axes[0,2].plot(t_smc, states_smc[:, 2], 'g-', label='SMC', linewidth=2)
    axes[0,2].plot(t_lqr, states_lqr[:, 2], 'b-', label='LQR', linewidth=2)
    axes[0,2].set_ylabel('Z Position (m)')
    axes[0,2].set_title('Z-Axis Position Tracking')
    axes[0,2].legend()
    axes[0,2].grid(True)
    
    # Position Error Norm
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    axes[1,0].plot(t_pid, error_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,0].plot(t_smc, error_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,0].plot(t_lqr, error_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,0].set_ylabel('Position Error Norm (m)')
    axes[1,0].set_xlabel('Time (s)')
    axes[1,0].set_title('Position Error Norm')
    axes[1,0].legend()
    axes[1,0].grid(True)
    
    # Control Input Norm
    control_norm_pid = np.linalg.norm(controls_pid, axis=1)
    control_norm_smc = np.linalg.norm(controls_smc, axis=1)
    control_norm_lqr = np.linalg.norm(controls_lqr, axis=1)
    
    axes[1,1].plot(t_pid, control_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,1].plot(t_smc, control_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,1].plot(t_lqr, control_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,1].set_ylabel('Control Input Norm')
    axes[1,1].set_xlabel('Time (s)')
    axes[1,1].set_title('Control Input Norm')
    axes[1,1].legend()
    axes[1,1].grid(True)
    
    # Disturbance profile
    disturbance_norm = np.linalg.norm(disturbances, axis=1)
    axes[1,2].plot(t_pid, disturbance_norm, 'k-', linewidth=2)
    axes[1,2].set_ylabel('Disturbance Magnitude')
    axes[1,2].set_xlabel('Time (s)')
    axes[1,2].set_title('Applied Disturbance Profile')
    axes[1,2].grid(True)
    
    # XY trajectory
    axes[2,0].plot(states_pid[:, 0], states_pid[:, 1], 'r-', label='PID', linewidth=2, alpha=0.7)
    axes[2,0].plot(states_smc[:, 0], states_smc[:, 1], 'g-', label='SMC', linewidth=2, alpha=0.7)
    axes[2,0].plot(states_lqr[:, 0], states_lqr[:, 1], 'b-', label='LQR', linewidth=2, alpha=0.7)
    axes[2,0].plot(x_ref, y_ref, 'k--', label='Desired', linewidth=2)
    axes[2,0].set_xlabel('X (m)')
    axes[2,0].set_ylabel('Y (m)')
    axes[2,0].set_title('XY Plane Trajectory')
    axes[2,0].legend()
    axes[2,0].grid(True)
    axes[2,0].axis('equal')
    
    # Performance metrics table
    metrics = calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, 
                                        controls_pid, controls_smc, controls_lqr,
                                        disturbance_type)
    
    controllers = ['PID', 'SMC', 'LQR']
    metric_names = list(metrics.keys())
    metric_values = list(metrics.values())
    
    # Create table
    cell_text = []
    for i, metric in enumerate(metric_names):
        row = [f'{metric_values[i][j]:.3f}' for j in range(3)]
        cell_text.append(row)
    
    table = axes[2,1].table(cellText=cell_text,
                          rowLabels=metric_names,
                          colLabels=controllers,
                          loc='center',
                          cellLoc='center')
    table.auto_set_font_size(False)
    table.set_fontsize(9)
    table.scale(1.2, 1.8)
    axes[2,1].axis('off')
    axes[2,1].set_title('Performance Metrics')
    
    # Recovery time comparison
    recovery_times = metrics['Recovery Time (s)']
    bars = axes[2,2].bar(controllers, recovery_times, color=['red', 'green', 'blue'], alpha=0.7)
    axes[2,2].set_ylabel('Recovery Time (s)')
    axes[2,2].set_title('Disturbance Recovery Time')
    axes[2,2].grid(True, alpha=0.3)
    
    for bar, time in zip(bars, recovery_times):
        axes[2,2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1, 
                      f'{time:.2f}s', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.show()
    
    return metrics

def main():
    """Main function to run all robustness tests"""
    
    # Test different disturbance types
    disturbance_types = [
        'constant',      # Constant external force
        'sinusoidal',    # Periodic disturbance
        'impulse',       # Sudden impacts
        'step',          # Sudden persistent disturbance
        'random',        # Random noise
        'complex'        # Combined disturbances
    ]
    
    all_results = {}
    
    for dist_type in disturbance_types:
        metrics = run_robustness_test(dist_type)
        all_results[dist_type] = metrics
        
        # Print results
        print(f"\nPerformance under {dist_type} disturbance:")
        print(f"{'Metric':<25} {'PID':<10} {'SMC':<10} {'LQR':<10}")
        print("-" * 60)
        for metric, values in metrics.items():
            print(f"{metric:<25} {values[0]:<10.3f} {values[1]:<10.3f} {values[2]:<10.3f}")
    
    # Summary comparison
    print(f"\n{'='*60}")
    print("OVERALL ROBUSTNESS SUMMARY")
    print(f"{'='*60}")
    
    # Calculate average performance across all tests
    avg_metrics = {}
    for metric in ['RMSE Position', 'Max Position Error', 'Steady-State Error', 
                  'Recovery Time (s)', 'Mean Control Effort', 'Max Control Effort']:
        avg_values = []
        for i in range(3):  # For each controller
            controller_avg = np.mean([all_results[dist][metric][i] for dist in disturbance_types])
            avg_values.append(controller_avg)
        avg_metrics[metric] = avg_values
    
    print("\nAverage Performance Across All Disturbance Tests:")
    print(f"{'Metric':<25} {'PID':<10} {'SMC':<10} {'LQR':<10}")
    print("-" * 60)
    for metric, values in avg_metrics.items():
        print(f"{metric:<25} {values[0]:<10.3f} {values[1]:<10.3f} {values[2]:<10.3f}")
    
    # Robustness ranking
    print(f"\n{'='*50}")
    print("ROBUSTNESS RANKING")
    print(f"{'='*50}")
    
    # Lower RMSE and recovery time is better, control effort should be moderate
    rmse_ranking = np.argsort(avg_metrics['RMSE Position'])
    recovery_ranking = np.argsort(avg_metrics['Recovery Time (s)'])
    control_ranking = np.argsort(avg_metrics['Mean Control Effort'])
    
    controllers = ['PID', 'SMC', 'LQR']
    
    print("Based on RMSE Position (Tracking Accuracy):")
    for i, rank in enumerate(rmse_ranking):
        print(f"{i+1}. {controllers[rank]} (RMSE: {avg_metrics['RMSE Position'][rank]:.3f}m)")
    
    print("\nBased on Recovery Time (Disturbance Rejection):")
    for i, rank in enumerate(recovery_ranking):
        print(f"{i+1}. {controllers[rank]} (Time: {avg_metrics['Recovery Time (s)'][rank]:.2f}s)")
    
    print("\nBased on Control Efficiency:")
    for i, rank in enumerate(control_ranking):
        print(f"{i+1}. {controllers[rank]} (Effort: {avg_metrics['Mean Control Effort'][rank]:.3f})")
    
    print("\nBased on Steady-State Performance:")
    steady_ranking = np.argsort(avg_metrics['Steady-State Error'])
    for i, rank in enumerate(steady_ranking):
        print(f"{i+1}. {controllers[rank]} (Error: {avg_metrics['Steady-State Error'][rank]:.3f}m)")

if __name__ == "__main__":
    main()